#!/bin/bash

# 无人机导航系统依赖安装脚本
# 适用于Ubuntu 22.04 + ROS2 Humble + Jetson Orin Nano

set -e  # 遇到错误立即退出

echo "=========================================="
echo "  无人机导航系统依赖安装脚本"
echo "  适用于 Jetson Orin Nano + ROS2 Humble"
echo "=========================================="

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

print_info() {
    echo -e "${BLUE}[信息]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[成功]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[警告]${NC} $1"
}

print_error() {
    echo -e "${RED}[错误]${NC} $1"
}

# 检查是否为root用户
check_root() {
    if [ "$EUID" -eq 0 ]; then
        print_error "请不要以root用户运行此脚本"
        exit 1
    fi
}

# 检查Ubuntu版本
check_ubuntu_version() {
    print_info "检查Ubuntu版本..."
    
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        if [ "$VERSION_ID" != "22.04" ] && [ "$VERSION_ID" != "20.04" ]; then
            print_warning "此脚本主要针对Ubuntu 22.04优化，当前版本: $VERSION_ID"
            read -p "是否继续安装? (y/n): " -n 1 -r
            echo
            if [[ ! $REPLY =~ ^[Yy]$ ]]; then
                exit 1
            fi
        fi
    fi
}

# 检查ROS2安装
check_ros2() {
    print_info "检查ROS2 Humble安装状态..."
    
    if [ -f /opt/ros/humble/setup.bash ]; then
        print_success "检测到ROS2 Humble已安装"
        source /opt/ros/humble/setup.bash
    else
        print_error "未检测到ROS2 Humble，请先安装ROS2"
        echo "安装指南: https://docs.ros.org/en/humble/Installation.html"
        exit 1
    fi
}

# 更新系统包
update_system() {
    print_info "更新系统包..."
    sudo apt update && sudo apt upgrade -y
    print_success "系统包更新完成"
}

# 安装基础依赖
install_basic_dependencies() {
    print_info "安装基础依赖包..."
    
    sudo apt install -y \
        build-essential \
        cmake \
        git \
        python3-pip \
        python3-dev \
        python3-setuptools \
        python3-wheel \
        libeigen3-dev \
        libgoogle-glog-dev \
        libgflags-dev \
        libatlas-base-dev \
        libsuitesparse-dev \
        curl \
        wget \
        unzip
        
    print_success "基础依赖安装完成"
}

# 安装ROS2开发工具
install_ros2_dev_tools() {
    print_info "安装ROS2开发工具..."
    
    sudo apt install -y \
        python3-colcon-common-extensions \
        python3-rosdep \
        python3-vcstool \
        python3-argcomplete \
        ros-humble-ament-cmake \
        ros-humble-ament-cmake-python
        
    # 初始化rosdep
    if [ ! -f /etc/ros/rosdep/sources.list.d/20-default.list ]; then
        print_info "初始化rosdep..."
        sudo rosdep init
    fi
    
    rosdep update
    print_success "ROS2开发工具安装完成"
}

# 安装Navigation2
install_navigation2() {
    print_info "安装Navigation2完整功能包..."
    
    sudo apt install -y \
        ros-humble-navigation2 \
        ros-humble-nav2-bringup \
        ros-humble-nav2-simple-commander \
        ros-humble-nav2-msgs \
        ros-humble-nav2-core \
        ros-humble-nav2-costmap-2d \
        ros-humble-nav2-planner \
        ros-humble-nav2-controller \
        ros-humble-nav2-recoveries \
        ros-humble-nav2-bt-navigator \
        ros-humble-nav2-waypoint-follower \
        ros-humble-nav2-lifecycle-manager
        
    print_success "Navigation2安装完成"
}

# 安装PCL和点云处理库
install_pcl_libraries() {
    print_info "安装PCL点云处理库..."
    
    sudo apt install -y \
        libpcl-dev \
        ros-humble-pcl-conversions \
        ros-humble-pcl-msgs \
        ros-humble-pcl-ros \
        ros-humble-perception-pcl
        
    print_success "PCL库安装完成"
}

# 安装传感器驱动
install_sensor_drivers() {
    print_info "安装传感器相关包..."
    
    sudo apt install -y \
        ros-humble-sensor-msgs \
        ros-humble-sensor-msgs-py \
        ros-humble-laser-geometry \
        ros-humble-laser-filters \
        ros-humble-pointcloud-to-laserscan
        
    print_success "传感器驱动安装完成"
}

# 安装TF和几何库
install_tf_geometry() {
    print_info "安装TF和几何变换库..."
    
    sudo apt install -y \
        ros-humble-tf2 \
        ros-humble-tf2-ros \
        ros-humble-tf2-geometry-msgs \
        ros-humble-tf2-sensor-msgs \
        ros-humble-tf2-tools \
        ros-humble-geometry-msgs \
        ros-humble-visualization-msgs
        
    print_success "TF和几何库安装完成"
}

# 安装Python依赖
install_python_dependencies() {
    print_info "安装Python依赖包..."
    
    pip3 install --user \
        numpy \
        scipy \
        matplotlib \
        scikit-learn \
        opencv-python \
        pyserial \
        pyyaml \
        transforms3d \
        numba
        
    print_success "Python依赖安装完成"
}

# 安装Livox SDK和驱动
install_livox_dependencies() {
    print_info "准备Livox相关依赖..."
    
    # 检查是否已有livox_ros_driver2
    if [ -d "src/livox_ros_driver2" ] || [ -d "../livox_ros_driver2" ]; then
        print_success "检测到Livox驱动已存在"
    else
        print_warning "Livox驱动需要手动安装，请参考README.md"
    fi
}

# 设置环境变量
setup_environment() {
    print_info "配置环境变量..."
    
    # 添加ROS2环境到bashrc
    if ! grep -q "source /opt/ros/humble/setup.bash" ~/.bashrc; then
        echo "source /opt/ros/humble/setup.bash" >> ~/.bashrc
        print_info "已添加ROS2环境到~/.bashrc"
    fi
    
    # 添加colcon自动补全
    if ! grep -q "source /usr/share/colcon_argcomplete/hook/colcon-argcomplete.bash" ~/.bashrc; then
        echo "source /usr/share/colcon_argcomplete/hook/colcon-argcomplete.bash" >> ~/.bashrc
        print_info "已添加colcon自动补全到~/.bashrc"
    fi
    
    print_success "环境变量配置完成"
}

# 优化Jetson性能
optimize_jetson_performance() {
    print_info "优化Jetson性能设置..."
    
    # 设置性能模式（如果是Jetson设备）
    if [ -f /usr/bin/jetson_clocks ]; then
        print_info "检测到Jetson设备，设置性能模式..."
        sudo jetson_clocks
        
        # 设置功耗模式
        if [ -f /usr/sbin/nvpmodel ]; then
            sudo nvpmodel -m 0  # 最大性能模式
            print_success "已设置为最大性能模式"
        fi
    fi
    
    # 优化系统设置
    echo "# Jetson优化设置" | sudo tee -a /etc/sysctl.conf
    echo "vm.swappiness=10" | sudo tee -a /etc/sysctl.conf
    echo "vm.vfs_cache_pressure=50" | sudo tee -a /etc/sysctl.conf
    
    print_success "Jetson性能优化完成"
}

# 验证安装
verify_installation() {
    print_info "验证安装状态..."
    
    source /opt/ros/humble/setup.bash
    
    # 检查关键包
    missing_packages=()
    
    if ! ros2 pkg list | grep -q "navigation2"; then
        missing_packages+=("navigation2")
    fi
    
    if ! ros2 pkg list | grep -q "nav2_simple_commander"; then
        missing_packages+=("nav2_simple_commander")
    fi
    
    if ! python3 -c "import numpy, scipy, cv2" 2>/dev/null; then
        missing_packages+=("python-libraries")
    fi
    
    if [ ${#missing_packages[@]} -eq 0 ]; then
        print_success "所有依赖验证通过！"
        return 0
    else
        print_error "以下包安装失败或缺失:"
        for pkg in "${missing_packages[@]}"; do
            echo "  - $pkg"
        done
        return 1
    fi
}

# 显示安装后说明
show_post_install_info() {
    echo ""
    echo "=========================================="
    echo -e "${GREEN}安装完成！${NC}"
    echo "=========================================="
    echo ""
    echo "下一步操作："
    echo "1. 重新打开终端或运行: source ~/.bashrc"
    echo "2. 进入工作空间: cd ~/drone_navigation_ws"
    echo "3. 编译项目: colcon build"
    echo "4. 安装环境: source install/setup.bash"
    echo "5. 运行测试: ros2 launch drone_navigation drone_navigation.launch.py"
    echo ""
    echo "重要提醒："
    echo "- 确保point_lio已正确安装和配置"
    echo "- 检查Livox Mid360驱动是否正常工作"
    echo "- 根据实际硬件调整配置文件参数"
    echo ""
    echo "如有问题，请查看README.md或联系开发团队"
}

# 主函数
main() {
    print_info "开始安装无人机导航系统依赖..."
    
    check_root
    check_ubuntu_version
    check_ros2
    
    update_system
    install_basic_dependencies
    install_ros2_dev_tools
    install_navigation2
    install_pcl_libraries
    install_sensor_drivers
    install_tf_geometry
    install_python_dependencies
    install_livox_dependencies
    setup_environment
    optimize_jetson_performance
    
    if verify_installation; then
        show_post_install_info
        exit 0
    else
        print_error "安装过程中出现问题，请检查错误信息"
        exit 1
    fi
}

# 检查参数
if [ "$1" = "--help" ] || [ "$1" = "-h" ]; then
    echo "无人机导航系统依赖安装脚本"
    echo ""
    echo "用法: $0 [选项]"
    echo ""
    echo "选项:"
    echo "  -h, --help     显示此帮助信息"
    echo "  --no-jetson    跳过Jetson优化"
    echo ""
    echo "环境要求:"
    echo "  - Ubuntu 22.04 (推荐) 或 20.04"
    echo "  - ROS2 Humble 已安装"
    echo "  - 8GB+ 内存 (Jetson Orin Nano)"
    echo ""
    exit 0
fi

# 跳过Jetson优化选项
if [ "$1" = "--no-jetson" ]; then
    optimize_jetson_performance() {
        print_info "跳过Jetson性能优化"
    }
fi

# 运行主函数
main 